package com.hlz.flink.chapter06.window;

import com.hlz.flink.chapter06.Stock;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.sql.Timestamp;
import java.time.Duration;
import java.util.HashSet;

/**
 * 全窗口函数
 *
 * @author Hongliang Zhu
 * @create 2023-06-23 20:39
 */
public class UvCountByWindowExample {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        DataStreamSource<String> dataStreamSource = env.socketTextStream("127.0.0.1", 9696);
        SingleOutputStreamOperator<Stock> stockStream = dataStreamSource.filter(StringUtils::isNotEmpty)
                .map(new MapFunction<String, Stock>() {
                    @Override
                    public Stock map(String value) throws Exception {
                        final String[] split = value.split(",");
                        return new Stock(split[0], Double.parseDouble(split[1]), split[2], Long.parseLong(split[3]));
                    }
                });

        // 设置水位线
        SingleOutputStreamOperator<Stock> stockSingleOutputStreamOperator1 = stockStream.assignTimestampsAndWatermarks(
                WatermarkStrategy.<Stock>forBoundedOutOfOrderness(Duration.ofMillis(2)).withTimestampAssigner(new SerializableTimestampAssigner<Stock>() {
                    @Override
                    public long extractTimestamp(Stock element, long recordTimestamp) {
                        return element.getTimeStamp();
                    }
                }));

        // 将数据全部发往同一分区，按窗口统计 UV

        stockSingleOutputStreamOperator1.keyBy(data -> true)
                .window(TumblingEventTimeWindows.of(Time.milliseconds(10)))
                .process( new UvCountByWindow())
                .print("window-function");
        env.execute();

    }

    public static class UvCountByWindow extends ProcessWindowFunction<Stock, String, Boolean, TimeWindow> {

        @Override
        public void process(Boolean aBoolean, Context context, Iterable<Stock> elements, Collector<String> out) throws Exception {
            HashSet<String> stockSet = new HashSet<>();
            for (Stock element : elements) {
                stockSet.add(element.getStockName());
            }
            // 结合窗口信息，包装输出内容
            long start = context.window().getStart();
            long end = context.window().getEnd();

            long watermark = context.currentWatermark();
            out.collect("窗口： " + new Timestamp(start) + "~" + new Timestamp(end) + "的独立访客数量是： " + stockSet.size() +", watermark is :" + watermark);
        }
    }
}
